from torch import nn
class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        self.lstm = nn.LSTM(input_size=1, hidden_size=6, num_layers=4, batch_first=True)
        self.line = nn.Linear(in_features=6, out_features=1)

    def forward(self, x):
        x, _ = self.lstm(x)
        x = x[:, -1, :]
        x = self.line(x)
        x = x.view(-1, 1, 1)
        return x